
import json
import os
from PIL import Image
#import imgaug.augmenters as iaa
import numpy as np
import labelme
from tqdm import  tqdm

#假设输入都是正方形

def dict_json(flags, imageData, shapes, imagePath, imageHeight,
              imageWidth):
    '''
    :param imageData: str
    :param shapes: list
    :param imagePath: str
    :return: dict
    '''
    return {"version": "4.5.6", "flags": flags, "shapes": shapes,
            'imagePath': imagePath, "imageData": imageData,
            'imageHeight': imageHeight,
            'imageWidth': imageWidth}


def resizeImageAndLabel(root,size):
    """
    整张图和标签（label）resize
    :param image: origin image
    :param pt: point
    :param size: size
    :return: new_image,new_pt
    """
    img_path = root + '/image'
    json_path = root + '/json'

    new_img_path= os.path.join(root,'new_image')
    new_json_path = os.path.join(root,'new_json')
    if not os.path.exists(new_img_path):
        os.makedirs(new_img_path)
    if not os.path.exists(new_json_path):
        os.makedirs(new_json_path)


    json_filenames = [json_filename for json_filename in os.listdir(json_path) if json_filename.endswith('.json')]
    for json_filename in tqdm(json_filenames):
        img_filename = json_filename.replace('.json', '.png')  # 单个图像的文件名
        jsonfile_path = os.path.join(json_path, json_filename)  # 完整的json路径
        imgfile_path = os.path.join(img_path, img_filename)  # 完整的图像路径

        image = Image.open(imgfile_path)  # 打开单个图像
        if image is None:
            continue

        j = open(jsonfile_path).read()  # json文件读入成字符串格式
        jj = json.loads(j)
        newshape = []

        origin_shape=image.width
        p=size/origin_shape #缩放倍率
        new_image=image.resize((size,size),Image.BICUBIC)

        for i in range(len(jj['shapes'])):
            pts = jj["shapes"][i]["points"]
            pts = np.array(pts).astype(int)
            pts[pts < 0] = 0
            pts=pts*p

            if len(pts) < 4:
                continue

            aa = {}
            aa.update({'label': jj["shapes"][i]["label"]})
            aa.update({"points": pts.tolist()})  # 此时的点是np格式 ，最终需要转换为列表格式
            aa.update({"group_id": None})
            aa.update({"shape_type": "polygon"})
            aa.update({"flags": {}})
            newshape.append(aa)
        flags = {}
        imageHeight = new_image.height
        imageWidth = new_image.width
        imagePath = json_filename[:-5]+'.png'
        imageData = None

        data_final = dict_json(flags, imageData, newshape, imagePath,
                               imageHeight, imageWidth)
        newfile_json = os.path.join(new_json_path, json_filename[:-5] + '.json')
        json.dump(data_final, open(newfile_json, 'w'), indent=4)
        newfile_image=os.path.join(new_img_path, json_filename[:-5] + '.png')
        new_image.save(newfile_image)
    return 0


def crop(root):
    new_image_dir_path=root+'/crop_image'
    if not os.path.exists(new_image_dir_path):
        os.makedirs(new_image_dir_path)
    image_ids=[image_id[:-4] for image_id in os.listdir(root) if image_id.endswith('.jpg')]
    for image_id in tqdm(image_ids):
        image_file_path=os.path.join(root,image_id+'.jpg')
        image=Image.open(image_file_path)
        origin_shape = image.width
        crop_pts=[[[0,0],[origin_shape/2,origin_shape/2]],
                  [[origin_shape/2,0],[origin_shape,origin_shape/2]],
                  [[0,origin_shape/2],[origin_shape/2,origin_shape]],
                  [[origin_shape/2,origin_shape/2],[origin_shape,origin_shape]]]

        for n,an_pt in enumerate (crop_pts):
            pt1=an_pt[0] #左上角
            pt2=an_pt[1] #右下角
            box=(pt1[0],pt1[1],pt2[0],pt2[1]) #(left,upper,right,lower) tuple
            new_image = image.crop(box)
            image_id=image_id+'-'+str(n)
            new_image_file_path=os.path.join(new_image_dir_path, image_id + '.jpg')
            new_image.save(new_image_file_path)
    return 0
def crop_sliding_window(root):
    #沿每个轴滑动窗口的次数 n ，一张图片的截取数目为 n*n
    new_crop_size = 512
    sliding_time = 4
    new_image_dir_path = root+ '/crop_image'
    if not os.path.exists(new_image_dir_path):
        os.makedirs(new_image_dir_path)
    image_ids = [image_id[:-4] for image_id in os.listdir(root) if image_id.endswith('.png')]
    for image_id in tqdm(image_ids):
        image_file_path = os.path.join(root, image_id + '.png')
        image = Image.open(image_file_path)
        sliding_distance=int((image.width-new_crop_size)/sliding_time)
        pt1 = np.array([0,0]).astype(int)
        pt2 =  np.array([new_crop_size,new_crop_size]).astype(int)

        for y in range(sliding_time):
            new_pt1 = pt1+[0,y*sliding_distance]
            new_pt2 = pt2+[0,y*sliding_distance]
            for x in range(sliding_time):
                new_pt1=new_pt1+[sliding_distance,0]
                new_pt2=new_pt2+[sliding_distance,0]
                box = (new_pt1[0], new_pt1[1], new_pt2[0], new_pt2[1])  # (left,upper,right,lower) tuple
                new_image = image.crop(box)
                new_image_id = image_id + '-x' + str(x)+'-y'+str(y)
                new_image_file_path = os.path.join(new_image_dir_path, new_image_id + '.png')
                new_image.save(new_image_file_path)


    return 0
def resize(root,size):
    new_image_dir_path = root + '/resize_image'
    if not os.path.exists(new_image_dir_path):
        os.makedirs(new_image_dir_path)
    image_ids = [image_id[:-4] for image_id in os.listdir(root) if image_id.endswith('.jpg')]
    for image_id in tqdm(image_ids):
        image_file_path=os.path.join(root,image_id+'.jpg')
        image=Image.open(image_file_path)
        new_image=image.resize((size,size),Image.BICUBIC)
        new_image_path=os.path.join(new_image_dir_path,image_id+'.jpg')
        new_image.save(new_image_path)

    return 0

def main():
    # root = r"F:\wyx\GanDatasets\large_data\large-2-single-size\sect-large\bigger"
    # resizeImageAndLabel(root,512)
    root=r'I:\wyx\data\Hik_zaishen\标注数据集\背景图'

    crop_sliding_window(root)

    #resizeImageAndLabel(root,1024)
    #resize(root,256)

if __name__ == '__main__':
    main()
